-
Notifications
You must be signed in to change notification settings - Fork 1.8k
MCPToolset: Add OAuth2 Client Credentials Flow with RFC 8414 Compliant Discovery #2061
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
@@ -307,6 +307,28 @@ | |||
* Added unit test coverage for local_eval_sets_manager.py ([174afb3](https://github.com/google/adk-python/commit/174afb3975bdc7e5f10c26f3eebb17d2efa0dd59)) | |||
* Extract common options for `adk web` and `adk api_server` ([01965bd](https://github.com/google/adk-python/commit/01965bdd74a9dbdb0ce91a924db8dee5961478b8)) | |||
|
|||
## [Unreleased] |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
CHANGELOG is only updated up on release. please put those information in the message.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
will do!
Thank you so much for the amazing PR @jeremyschulman and I do think this feature is useful !! However such PR is too big to review, would you please kindly split it into multiple smaller PRs (group those logically highly related changes into one PR) , thank you :) |
@seanzhou1023 - Could you please provide guidance on how you would like this PR split into parts? From a Developer perspective, the feature enhancement is all part-and-parcel together. Happy to split it up, and I value your help. |
Hey @seanzhou1023 I'm working with Jeremy on this and wanted to propose this and get input for how to split.
I know we could put 3 inside 1 and 2 partially but we're trying to be mindful on sizing of the PRs. |
Never mind, I'm in the middle of review, it just takes longer to review big PRs. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Overall review comments:
- Vibe coding is a good tool, but there are many traps, let's review LLM-generated codes carefully.
- Let's focus on implementing RFC 8414 in the PR, remove other unrelated fixes / features from this PR to keep it clean and clear given it's already very big. (e.g. verify_ssl & client_credentials grant type support is not specific to RFC 8414)
- logs are too much , it's making codes fragile and hard to read. also let's remove all emoji from the logs
@@ -12,11 +12,39 @@ | |||
# See the License for the specific language governing permissions and | |||
# limitations under the License. | |||
|
|||
"""Auth configurations for Google ADK.""" |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
let's focus on the feature implementation in the PR and don't bother to manage the exported symbols in the package. I will fix them uniformly. Le't revert the changes of this file so that changes in this PR is more focused and and change history is more clear about what's done in this PR given this PR is already big enough.
@@ -50,16 +50,32 @@ class OAuthGrantType(str, Enum): | |||
PASSWORD = "password" | |||
|
|||
@staticmethod | |||
def from_flow(flow: OAuthFlows) -> "OAuthGrantType": | |||
"""Converts an OAuthFlows object to a OAuthGrantType.""" | |||
def from_flow(flow: OAuthFlows) -> Optional["OAuthGrantType"]: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
similarly for this file, let's don't bother to fix this file and focus on implementing the feature itself. I will create a separate PR to fix this file. let's revert the changes of this file so that changes in this PR is more focused and change history is more clear about what's done in this PR given this PR is already big enough.
await self._validate_credential() | ||
logger.debug("✅ Step 1: Credential validation passed") |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
can we remove those emoji from the log, we don't have emoji in other logs , this makes the logs inconsistent. you can probably instruct the LLM not to use emoji
The prepared authentication credential, or None if unavailable. | ||
""" | ||
|
||
logger.debug("🔄 CredentialManager.get_auth_credential() called") |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I feel the logs are too verbose in this file, it's kind of tracing every line of codes which makes the codes quite fragile and hard to read especially for those if-else
branch specifically for logging only. TBH, with those logging codes, I cannot easily figure out what the actual intention of the changes in this file is. Let's only log critical information. you could probably instruct the LLM not to log too much.
if self._auth_config.exchanged_auth_credential: | ||
return self._auth_config.exchanged_auth_credential | ||
# Then try to load from context | ||
if hasattr(callback_context, "_auth_credential") and callback_context._auth_credential: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
when will we have _auth_credential in callback_context ? storing credential in context is risky, is it hallucinated by LLM ?
base_url = f"{parsed.scheme}://{parsed.netloc}" | ||
logger.debug(f"Auto-detected OAuth discovery base URL: {base_url} from MCP URL: {full_url}") | ||
|
||
elif isinstance(self._connection_params, SseConnectionParams): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
why bother having separate elif for SseConnectionParams given the logic is same as StreamableHTTPConnectionParams ?
else: | ||
# For Stdio connections, OAuth discovery is not applicable | ||
logger.debug("❌ Disabling OAuth discovery for non-HTTP connection (Stdio)") | ||
return MCPAuthDiscovery( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
will None be enable to indicate AuthDiscovery is not applicable ?
elif isinstance(self._auth_scheme, OAuth2): | ||
# Check if OAuth2 scheme has client credentials flow with empty/invalid tokenUrl | ||
if (self._auth_scheme.flows and | ||
self._auth_scheme.flows.clientCredentials and |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think RFC 8414 is not specific to client_credentials grant type ?
base_url = self._auth_discovery.base_url | ||
logger.debug(f"Using explicitly configured discovery base URL: {base_url}") | ||
else: | ||
# Auto-extract base URL from connection parameters (same logic as _create_default_auth_discovery) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
given they are same, can we extract the common logic to private functions ?
discovery_scopes = None | ||
if (isinstance(self._auth_scheme, OAuth2) and | ||
self._auth_scheme.flows and | ||
self._auth_scheme.flows.clientCredentials and |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
why only handle client_credentials ? RFC 8414 is not specific to client_credentials.
also , please make sure it's not breaking existing functionalities. I tested the example |
dummy_context = DummyCallbackContext() | ||
|
||
try: | ||
# Exchange credentials to get access token with SSL verification setting |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don't think it's a good idea to do credential exchange here , given MCP tool is already capable of exchanging access tokens.
# Get session from session manager | ||
session = await self._mcp_session_manager.create_session() | ||
# Perform OAuth discovery if needed | ||
await self._perform_oauth_discovery() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
given oauth discovery is not specific to MCP toolset, should we put such logic in credential manager ?
""" | ||
# For oauth-protected-resource discovery | ||
if "authorization_servers" in config: | ||
return isinstance(config["authorization_servers"], list) and bool(config["authorization_servers"]) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
do we just want to validate config["authorization_servers"] is a non-empty list here ? would
isinstance(config["authorization_servers"], list) and len(config["authorization_servers"]) > 0
more readable ?
@experimental | ||
def _validate_oauth_discovery_response(config: Dict[str, Any]) -> bool: | ||
""" | ||
Validate OAuth discovery response contains required fields. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
better document what's validated here and why (i.e. what's expected response from different discovery endpoint)
True if configuration is valid, False otherwise | ||
""" | ||
# For oauth-protected-resource discovery | ||
if "authorization_servers" in config: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
such check is not reliable, e.g. returning authorization_servers
in oauth-authorization-server
is not valid. better explicitly pass in type of discovery to this method or have separate validation method to do the validation for different discovery endpoint.
Create an OAuth2 auth scheme by automatically discovering OAuth configuration. | ||
|
||
Implements RFC 8414 two-stage discovery: | ||
1. Query .well-known/oauth-protected-resource to find authorization server |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
is this part of RFC 8414 ? I don't see it : https://datatracker.ietf.org/doc/html/rfc8414
|
||
token_endpoint = None | ||
|
||
if protected_resource_config and "authorization_servers" in protected_resource_config: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
do we need to check "authorization_servers" in protected_resource_config again ? I think we already validate it in _validate_oauth_discovery_response
auth_server_url, OAUTH_AUTHORIZATION_SERVER_DISCOVERY, timeout, verify_ssl | ||
) | ||
|
||
if auth_server_config and "token_endpoint" in auth_server_config: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
do we need to check "token_endpoint" in auth_server_config again ? I think we already validate it in _validate_oauth_discovery_response
|
||
|
||
@experimental | ||
async def create_oauth_scheme_from_discovery( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
do we need to manually implement it from scratch(sending the http request and parse the result ourselves) , aren't there existing library to leverage given it's a standard ?
grant_type = self._get_grant_type(auth_scheme) | ||
logger.debug(f"🎯 detected grant type: {grant_type}") | ||
|
||
if grant_type == OAuthGrantType.CLIENT_CREDENTIALS: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Client_Credential flow is othorgonal to RFC 8414 , let's separate them into different PRs.
Thanks for this pull request. One thing we may wanna do in this PR, or in a followup PR is to also try out OpenID Connect Discovery 1.0 endpoints after OAuth 2.0 Authorization Server Metadata endpoints as mentioned in MCP authorization spec: |
Add OAuth2 Client Credentials Flow with Automatic Discovery to ADK MCPToolset
The ADK now offers OAuth2 client credentials authentication that "just works" out of the box! 🚀
Summary
This PR adds enterprise-grade OAuth2 client credentials authentication to the ADK MCPToolset with automatic discovery capabilities. The implementation follows RFC 8414 for OAuth2 Authorization Server Metadata discovery and provides seamless authentication setup with minimal configuration.
No breaking changes - This PR is fully backwards compatible. Existing MCPToolset usage continues to work unchanged.
Key Features
🚀 Automatic OAuth Discovery by Default
🔧 RFC 8414 Compliant Discovery
.well-known/oauth-protected-resource
to find authorization server.well-known/oauth-authorization-server
for token endpoint🔐 Complete OAuth2 Client Credentials Flow
📝 Production-Ready Logging
Files Updated
🆕 New Files Created
adk-python/src/google/adk/auth/oauth2_discovery_util.py
(~180 lines)Purpose: OAuth2 discovery utilities following RFC 8414
Key Functions:
create_oauth_scheme_from_discovery()
- Main discovery entry point_query_oauth_protected_resource()
- Query .well-known/oauth-protected-resource_query_authorization_server_metadata()
- Query authorization server metadata_create_oauth2_scheme()
- Create OAuth2 scheme from discovered endpointsadk-python/src/google/adk/tools/mcp_tool/mcp_auth_discovery.py
(~85 lines)Purpose: Configuration class for OAuth discovery
Key Features:
MCPAuthDiscovery
dataclass with base_url, timeout, enabled properties📝 Enhanced Existing Files
adk-python/src/google/adk/tools/mcp_tool/mcp_toolset.py
(~100 lines added/modified)Major Changes:
MCPAuthDiscovery
from connection params_create_default_auth_discovery()
: Extracts base URLs from HTTP connections_perform_oauth_discovery()
: Performs RFC 8414 two-stage discoveryauth_discovery
parameter with automatic defaultsadk-python/src/google/adk/auth/credential_manager.py
(~50 lines added/modified)Key Enhancements:
_exchanger_registry
get_auth_credential()
: 8-step credential processing workflowadk-python/src/google/adk/auth/exchanger/oauth2_credential_exchanger.py
(~80 lines added/modified)Major Enhancements:
_exchange_client_credentials()
methodclient_secret_post
method_get_grant_type()
for flow type determinationadk-python/src/google/adk/tools/mcp_tool/mcp_tool.py
(~10 lines added/modified)Minor Enhancements:
_get_headers()
for OAuth2 tokensUsage Examples
Automatic Discovery
Override Options
Integration Flow
The implementation follows a clean integration pattern:
MCPToolset
→ Creates defaultMCPAuthDiscovery
from connection paramsoauth2_discovery_util
→ Discovers OAuth endpoints via RFC 8414credential_manager
→ Registers and usesOAuth2CredentialExchanger
oauth2_credential_exchanger
→ Performs client credentials token exchangemcp_tool
→ Uses exchanged tokens for authenticated MCP requestsDefault Behavior Logic
Auto-Enable OAuth Discovery:
http://localhost:9204/mcp/
→http://localhost:9204
(discovery base)http://server:8080/sse/
→http://server:8080
(discovery base)Discovery Process:
.well-known/oauth-protected-resource
for authorization server.well-known/oauth-authorization-server
for token endpointBenefits
auth_discovery
parameters still workTechnical Details
OAuth2 Discovery Implementation
Authentication Flow
Logging Strategy
Code Statistics
Test Coverage
🧪 Comprehensive Test Suite Added
The OAuth2 enhancement includes comprehensive test coverage across all new functionality:
New Test Files Created
test_mcp_auth_discovery.py
(~140 lines)test_mcp_toolset_oauth_discovery.py
(~380 lines)test_credential_manager_oauth2_integration.py
(~140 lines)Enhanced Existing Test Coverage
Existing comprehensive tests already cover:
test_oauth2_discovery_util.py
- OAuth2 discovery utilities (RFC 8414)test_oauth2_credential_exchanger.py
- Client credentials flow implementationtest_credential_manager.py
- Credential management workflows🎯 Test Coverage Areas
📊 Test Statistics
Dependencies
Sample Implementation
🎯 Comprehensive OAuth2 Client Credentials Sample
A complete sample implementation has been created at:
adk-python/contributing/samples/mcp_oauth2_client_credentials_agent/
Sample Features:
Sample Contents:
agent.py
- Five different OAuth2 agent configurationsmock_oauth_server.py
- Complete OAuth2 test server implementationREADME.md
- Comprehensive documentation and usage guide__init__.py
- Standard Python package structureDemonstrated Scenarios:
Mock OAuth2 Server Features:
.well-known/oauth-protected-resource
,.well-known/oauth-authorization-server
)